"""
@File    :  arp_dep_test_layout.py
@Author  :  LeeSen
@Date    :  2024-07-12
@Version :  Python 3.7.4
"""
import time
import IPy
from PyQt5.QtCore import QThread, pyqtSignal
from PyQt5.QtWidgets import QVBoxLayout, QGridLayout, QHBoxLayout, QLabel, \
    QTextEdit, QPushButton, QComboBox, QLineEdit, QMessageBox
from scapy.sendrecv import sniff
from common_config_param import GlobalParam
from common_function import welcome_init, ip_format_check, int_num_check, validate_mac, init_nic, log_format_print, \
    clear_log
from scapy_function import generate_arp_packet, my_send_pkt, send_icmp_request


def int_to_mac(num):
    """模拟Device数量生成MAC"""
    a = int(num / 65536)
    b = num % 65536
    c = int(b / 256)
    d = b % 256
    return "04:00:00:" + hex(a)[2:4].zfill(2) + ":" + hex(c)[2:4].zfill(2) + ":" + hex(d)[2:4].zfill(2)


class ArpSniffThread(QThread):
    pkt_signal = pyqtSignal(dict)

    def __init__(self, config_data):
        super().__init__()
        self.stop_flag = False
        self.filter_des = "arp"
        self.config_data = config_data

    def run(self):
        res = {"code": "200", "message": "Sniff Thread will Start", "data": ""}
        self.pkt_signal.emit(res)
        try:
            sniff(filter=self.filter_des,
                  iface=self.config_data["networkAdapterDescription"],
                  stop_filter=lambda x: self.stop_flag,
                  prn=lambda pkt: self.package_handler(pkt, self.config_data, res))
        except:
            res["code"] = "404"
            res["message"] = "Sniff Thread have Except"
            self.pkt_signal.emit(res)
        res["code"] = "200"
        res["message"] = "Sniff Thread will Stop"
        self.pkt_signal.emit(res)

    def stop(self):
        self.stop_flag = True

    def package_handler(self, pkt, config_data, res):
        item = pkt[0]
        if item.op == 1 and item.src == config_data["gatewayMAC"]:
            if self.config_data["receiveArpRequestPrint"]:
                res["code"] = "666"
                res["message"] = "Receive ARP Request:"
                res["data"] = item['Ether'].mysummary() + " / " + item['ARP'].mysummary()
                self.pkt_signal.emit(res)
            ipStr = IPy.IP(item.pdst)
            numID = (ipStr.int()) - (IPy.IP(config_data["deviceStartIP"]).int())
            usedMac = int_to_mac(numID)
            arp_pkt = generate_arp_packet(item.hwsrc,
                                          usedMac,
                                          usedMac,
                                          item.pdst,
                                          item.hwsrc,
                                          item.psrc, op_code=2)
            my_send_pkt(arp_pkt, self.config_data["networkAdapterDescription"])
            if self.config_data["sendArpReplyPrint"]:
                res["code"] = "666"
                res["message"] = "Send ARP Reply:"
                res["data"] = arp_pkt['Ether'].mysummary() + " / " + arp_pkt['ARP'].mysummary()
                self.pkt_signal.emit(res)
            if self.config_data["sendArpRelpyAfterSendIcmpRequest"]:
                send_icmp_request(item.pdst,
                                  usedMac,
                                  item.psrc,
                                  item.hwsrc,
                                  self.config_data["networkAdapterDescription"])

        elif item.op == 2 and item.src == config_data["gatewayMAC"]:
            if self.config_data["receiveArpReplyPrint"]:
                res["code"] = "666"
                res["message"] = "Receive ARP Reply:"
                res["data"] = item['Ether'].mysummary() + " / " + item['ARP'].mysummary()
                self.pkt_signal.emit(res)
            if self.config_data["receivedArpReplyAfterSendIcmpRequest"]:
                send_icmp_request(item.pdst,
                                  item.hwdst,
                                  item.psrc,
                                  item.hwsrc,
                                  self.config_data["networkAdapterDescription"])


class ArpDepTestThread(QThread):
    pkt_signal = pyqtSignal(dict)

    def __init__(self, config_data):
        super().__init__()

        self.stop_flag = False
        self.config_data = config_data

    def run(self):
        res = {"code": "200", "message": "ARP Dep Test Thread will Start", "data": ""}
        self.pkt_signal.emit(res)
        while not self.stop_flag:
            for i in range(0, self.config_data["deviceNum"]):
                used_ip = IPy.intToIp((IPy.IP(self.config_data["deviceStartIP"]).int() + i), 4)
                used_mac = int_to_mac(i)
                arp_pkt = generate_arp_packet("ff:ff:ff:ff:ff:ff",
                                              used_mac,
                                              used_mac,
                                              used_ip,
                                              "00:00:00:00:00:00",
                                              self.config_data["gatewayIP"], 1)
                my_send_pkt(arp_pkt, self.config_data["networkAdapterDescription"])
                if self.config_data["sendArpRequestPrint"]:
                    res["code"] = "666"
                    res["message"] = "Send ARP Request:"
                    res["data"] = arp_pkt['Ether'].mysummary() + " / " + arp_pkt['ARP'].mysummary()
                    self.pkt_signal.emit(res)
                time.sleep(self.config_data["singleArpInterval"])
                if self.stop_flag:
                    break
        res["code"] = "200"
        res["message"] = "ARP Dep Test Thread will Stop"
        self.pkt_signal.emit(res)

    def stop(self):
        self.stop_flag = True


class ArpDepTestLayout(QVBoxLayout):
    def __init__(self):
        super().__init__()

        self.thread_dict = {}
        self.config_data = {"deviceNum": 10,
                            "deviceStartIP": "1.1.1.2",
                            "gatewayIP": "1.1.1.1",
                            "gatewayMAC": "00:00:00:aa:bb:cc",
                            "networkAdapterDescription": "Realtek PCIe GBE Family Controller",
                            "singleArpInterval": 1,
                            "sendArpRelpyAfterSendIcmpRequest": True,
                            "receivedArpReplyAfterSendIcmpRequest": True,
                            "sendArpRequestPrint": False,
                            "sendArpReplyPrint": False,
                            "receiveArpRequestPrint": False,
                            "receiveArpReplyPrint": False
                            }

        top_layout = QHBoxLayout()

        config_param_layout = QGridLayout()

        nic_label = QLabel("Network Adapter:")
        self.nic_config_box = QComboBox()

        self.start_button = QPushButton('Start')
        self.start_button.clicked.connect(self.start_test)
        self.start_button.setEnabled(True)
        self.stop_button = QPushButton('Stop')
        self.stop_button.clicked.connect(self.stop_test)
        self.stop_button.setEnabled(False)

        self.clear_log_button = QPushButton('Clear Log')
        self.clear_log_button.clicked.connect(lambda x: clear_log(self.log_area))

        self.reload_network_adapter_btn = QPushButton("Reload NIC")
        self.reload_network_adapter_btn.clicked.connect(lambda x: init_nic(self.nic_config_box))

        config_param_layout.addWidget(nic_label, 0, 0)
        config_param_layout.addWidget(self.nic_config_box, 0, 1)
        config_param_layout.addWidget(self.start_button, 0, 3)
        config_param_layout.addWidget(self.stop_button, 0, 4)
        config_param_layout.addWidget(self.clear_log_button, 0, 5)
        config_param_layout.addWidget(self.reload_network_adapter_btn, 0, 6)

        mode_label = QLabel("Mode:")
        self.mode_config_box = QComboBox()
        self.mode_config_box.addItem("Keep Online Host")

        device_num_label = QLabel("Device Num:")
        self.device_num = QComboBox()

        starting_ip_label = QLabel("Starting IP:")
        self.starting_ip_input = QLineEdit()
        self.starting_ip_input.setText("192.168.0.2")

        arp_interval_label = QLabel("ARP Send Interval:[1-60]")
        self.arp_interval_input = QLineEdit()
        self.arp_interval_input.setText("1")

        gateway_ip_label = QLabel("Gateway IP:")
        self.gateway_ip_input = QLineEdit()
        self.gateway_ip_input.setText("192.168.0.1")

        gateway_mac_label = QLabel("Gateway MAC:")
        self.gateway_mac_input = QLineEdit()
        self.gateway_mac_input.setText("00:1e:cd:00:00:01")

        config_param_layout.addWidget(mode_label, 1, 0)
        config_param_layout.addWidget(self.mode_config_box, 1, 1)
        config_param_layout.addWidget(device_num_label, 2, 0)
        config_param_layout.addWidget(self.device_num, 2, 1)
        config_param_layout.addWidget(starting_ip_label, 3, 0)
        config_param_layout.addWidget(self.starting_ip_input, 3, 1)
        config_param_layout.addWidget(arp_interval_label, 4, 0)
        config_param_layout.addWidget(self.arp_interval_input, 4, 1)
        config_param_layout.addWidget(gateway_ip_label, 5, 0)
        config_param_layout.addWidget(self.gateway_ip_input, 5, 1)
        config_param_layout.addWidget(gateway_mac_label, 6, 0)
        config_param_layout.addWidget(self.gateway_mac_input, 6, 1)

        sendArpRelpyAfterSendIcmpRequest_label = QLabel("Send ICMP Request When Send ARP Relpy:")
        self.sendArpRelpyAfterSendIcmpRequest_box = QComboBox()
        self.sendArpRelpyAfterSendIcmpRequest_box.addItem("Yes")
        self.sendArpRelpyAfterSendIcmpRequest_box.addItem("No")

        receivedArpReplyAfterSendIcmpRequest_label = QLabel("Send ICMP Request When Received ARP Reply:")
        self.receivedArpReplyAfterSendIcmpRequest_box = QComboBox()
        self.receivedArpReplyAfterSendIcmpRequest_box.addItem("Yes")
        self.receivedArpReplyAfterSendIcmpRequest_box.addItem("No")

        send_arp_request_pkt_print_switch_label = QLabel("Send ARP Request pkt Print:")
        self.send_arp_request_pkt_print_switch_box = QComboBox()
        self.send_arp_request_pkt_print_switch_box.addItem("Off")
        self.send_arp_request_pkt_print_switch_box.addItem("On")

        send_arp_reply_pkt_print_switch_label = QLabel("Send ARP Reply pkt Print:")
        self.send_arp_reply_pkt_print_switch_box = QComboBox()
        self.send_arp_reply_pkt_print_switch_box.addItem("Off")
        self.send_arp_reply_pkt_print_switch_box.addItem("On")

        receive_arp_request_pkt_print_switch_label = QLabel("Receive ARP Request pkt Print:")
        self.receive_arp_request_pkt_print_switch_box = QComboBox()
        self.receive_arp_request_pkt_print_switch_box.addItem("Off")
        self.receive_arp_request_pkt_print_switch_box.addItem("On")

        receive_arp_reply_pkt_print_switch_label = QLabel("Receive ARP Reply pkt Print:")
        self.receive_arp_reply_pkt_print_switch_box = QComboBox()
        self.receive_arp_reply_pkt_print_switch_box.addItem("Off")
        self.receive_arp_reply_pkt_print_switch_box.addItem("On")

        config_param_layout.addWidget(sendArpRelpyAfterSendIcmpRequest_label, 1, 2)
        config_param_layout.addWidget(self.sendArpRelpyAfterSendIcmpRequest_box, 1, 3)

        config_param_layout.addWidget(receivedArpReplyAfterSendIcmpRequest_label, 2, 2)
        config_param_layout.addWidget(self.receivedArpReplyAfterSendIcmpRequest_box, 2, 3)

        config_param_layout.addWidget(send_arp_request_pkt_print_switch_label, 3, 2)
        config_param_layout.addWidget(self.send_arp_request_pkt_print_switch_box, 3, 3)

        config_param_layout.addWidget(send_arp_reply_pkt_print_switch_label, 4, 2)
        config_param_layout.addWidget(self.send_arp_reply_pkt_print_switch_box, 4, 3)

        config_param_layout.addWidget(receive_arp_request_pkt_print_switch_label, 5, 2)
        config_param_layout.addWidget(self.receive_arp_request_pkt_print_switch_box, 5, 3)

        config_param_layout.addWidget(receive_arp_reply_pkt_print_switch_label, 6, 2)
        config_param_layout.addWidget(self.receive_arp_reply_pkt_print_switch_box, 6, 3)

        config_param_layout.setContentsMargins(10, 10, 10, 10)
        config_param_layout.setSpacing(10)

        top_layout.addLayout(config_param_layout)

        self.addLayout(top_layout)

        log_label = QLabel("Real-time Logs")
        self.log_area = QTextEdit()
        self.log_area.setPlaceholderText("Log区域")
        self.log_area.setReadOnly(True)

        self.addWidget(log_label)
        self.addWidget(self.log_area)

        welcome_init(self.log_area)

        init_nic(self.nic_config_box)
        self.init_device_num()

    def init_device_num(self):
        for i in range(1, 16):
            self.device_num.addItem(str(i))
        for i in range(16, 10241, 16):
            self.device_num.addItem(str(i))

    def start_test(self):
        if GlobalParam.global_running_flag:
            QMessageBox.warning(self.parentWidget().parentWidget(), "WARNING",
                                "Other programs are running, please stop first")
            return
        if self.config_param_check():
            return
        self.update_config_data()
        sniffItem = ArpSniffThread(self.config_data)
        self.thread_dict["sniff"] = sniffItem
        sniffItem.pkt_signal.connect(self.update_message_or_data)
        self.thread_dict["sniff"].start()
        arpDepThread = ArpDepTestThread(self.config_data)
        self.thread_dict["arpDep"] = arpDepThread
        arpDepThread.pkt_signal.connect(self.update_message_or_data)
        self.thread_dict["arpDep"].start()
        self.start_button.setEnabled(False)
        self.stop_button.setEnabled(True)
        GlobalParam.global_running_flag = True

    def update_message_or_data(self, res):
        if res["code"] == "404":
            self.stop_test()
            log_format_print(self.log_area, res["message"])
        elif res["code"] == "200":
            log_format_print(self.log_area, res["message"])
        elif res["code"] == "666":
            log_format_print(self.log_area, res["message"] + res["data"])

    def stop_test(self):
        while self.thread_dict["sniff"].isRunning():
            self.thread_dict["sniff"].stop()
        while self.thread_dict["arpDep"].isRunning():
            self.thread_dict["arpDep"].stop()
        self.start_button.setEnabled(True)
        self.stop_button.setEnabled(False)
        self.thread_dict.pop("sniff")
        self.thread_dict.pop("arpDep")
        GlobalParam.global_running_flag = False

    def config_param_check(self):
        if ip_format_check(self.starting_ip_input.text()):
            QMessageBox.warning(self.parentWidget().parentWidget(),
                                "Invalid IP",
                                "The entered Starting IP address is invalid.")
            return True
        if ip_format_check(self.gateway_ip_input.text()):
            QMessageBox.warning(self.parentWidget().parentWidget(),
                                "Invalid IP",
                                "The entered Gateway IP address is invalid.")
            return True
        arp_interval = self.arp_interval_input.text()
        if int_num_check(arp_interval) or (int(arp_interval) < 1 or int(arp_interval) > 60):
            QMessageBox.warning(self.parentWidget().parentWidget(),
                                "Invalid ARP interval Param",
                                "The entered ARP interval Param is invalid.[1-60]")
            return True
        gateway_mac = self.gateway_mac_input.text()
        if not validate_mac(gateway_mac):
            QMessageBox.warning(self.parentWidget().parentWidget(),
                                "Invalid MAC",
                                "The entered Gateway MAC address is invalid.")
            return True

        return False

    def update_config_data(self):
        self.config_data["deviceNum"] = int(self.device_num.currentText())
        self.config_data["deviceStartIP"] = self.starting_ip_input.text()
        self.config_data["gatewayIP"] = self.gateway_ip_input.text()
        self.config_data["gatewayMAC"] = (self.gateway_mac_input.text()).replace("-", ":").lower()
        self.config_data["networkAdapterDescription"] = self.nic_config_box.currentText()
        self.config_data["singleArpInterval"] = int(self.arp_interval_input.text())
        self.config_data[
            "sendArpRelpyAfterSendIcmpRequest"] = True if self.sendArpRelpyAfterSendIcmpRequest_box.currentText() == "Yes" else False
        self.config_data[
            "receivedArpReplyAfterSendIcmpRequest"] = True if self.receivedArpReplyAfterSendIcmpRequest_box.currentText() == "Yes" else False
        self.config_data[
            "sendArpRequestPrint"] = True if self.send_arp_request_pkt_print_switch_box.currentText() == "On" else False
        self.config_data[
            "sendArpReplyPrint"] = True if self.send_arp_reply_pkt_print_switch_box.currentText() == "On" else False
        self.config_data[
            "receiveArpRequestPrint"] = True if self.receive_arp_request_pkt_print_switch_box.currentText() == "On" else False
        self.config_data[
            "receiveArpReplyPrint"] = True if self.receive_arp_reply_pkt_print_switch_box.currentText() == "On" else False

